import os
import asyncio
import json
import time
import datetime
from playwright.async_api import async_playwright


class LoginManager:
    def __init__(self, cookies_path='cookies.json'):
        self.cookies_path = cookies_path
        self.browser = None
        self.context = None
        self.page = None

    async def load_cookies(self):
        """加载 cookies.json 并修复格式问题"""
        if os.path.exists(self.cookies_path):
            with open(self.cookies_path, 'r', encoding='utf-8') as f:
                cookies = json.load(f)
                for cookie in cookies:
                    if cookie.get('domain') == 'www.douyin.com':
                        cookie['domain'] = '.douyin.com'
                    if isinstance(cookie.get('expires'), str):
                        try:
                            dt = datetime.datetime.fromisoformat(cookie['expires'].replace("Z", "+00:00"))
                            cookie['expires'] = int(dt.timestamp())
                        except Exception:
                            cookie['expires'] = None
                    cookie.setdefault('httpOnly', False)
                    cookie.setdefault('secure', True)
                    if cookie.get('sameSite') not in ['Strict', 'Lax', 'None']:
                        cookie['sameSite'] = 'None'
                return cookies
        return []

    async def save_cookies(self):
        """保存当前上下文中的 cookies 到文件"""
        if self.context:
            cookies = await self.context.cookies()
            with open(self.cookies_path, 'w', encoding='utf-8') as f:
                json.dump(cookies, f, indent=2, ensure_ascii=False)
            print("🍪 登录态已保存到 cookies.json")

    async def login(self):
        """启动浏览器并执行登录抖音流程"""
        playwright = await async_playwright().start()
        try:
            self.browser = await playwright.chromium.launch(headless=False, args=[
                '--disable-blink-features=AutomationControlled',
                '--disable-infobars',
                '--disable-automation',
                '--no-sandbox',
                '--disable-setuid-sandbox',
            ])
            self.context = await self.browser.new_context(
                accept_downloads=True,
                bypass_csp=True,
            )

            # 加载本地 Cookie（如果存在）
            cookies = await self.load_cookies()
            if cookies:
                await self.context.add_cookies(cookies)
                print("✅ 已加载本地登录态")

            # 打开页面
            self.page = await self.context.new_page()
            await self.page.add_init_script("""
                delete navigator.__proto__.webdriver;
                window.chrome = {
                    runtime: {},
                    csi: function() { return { onloadTime: new Date().getTime() }; },
                    loadTimes: function() {
                        return { requestTime: 0 };
                    }
                };
                navigator.chrome = {
                    app: { isInstalledNaClPlugin: false },
                    webstore: { flexibleName: true },
                    runtime: {}
                };
            """)
            await self.page.goto("https://www.douyin.com")

            start_time = time.time()
            timeout = 8 * 60  # 8 分钟
            login_prompt_locator = self.page.locator("div > div > div.QgsIceRc > div")

            await asyncio.sleep(2)
            while time.time() - start_time < timeout:
                if await login_prompt_locator.is_visible():
                    print("⏳ 当前尚未登录，继续等待...")
                else:
                    print("✅ 已登录")
                    break
                await asyncio.sleep(5)

            # 保存 Cookie
            await self.save_cookies()

            return True
        except Exception as e:
            print(f"登录过程中发生错误: {e}")
            raise
        finally:
            # 确保资源被正确关闭
            await self.close()
            await playwright.stop()

    async def close(self):
        """关闭浏览器资源"""
        if self.page:
            await self.page.close()
            self.page = None
        if self.context:
            await self.context.close()
            self.context = None
        if self.browser:
            await self.browser.close()
            self.browser = None


# 测试入口
if __name__ == "__main__":
    async def main():
        login_manager = LoginManager()
        try:
            await login_manager.login()
        except Exception as e:
            print(f"程序运行出错: {e}")
        finally:
            # 确保所有资源都被清理
            await login_manager.close()

    asyncio.run(main())